# -*- coding: utf-8 -*-
'''
Created on 2017年5月22日

@author: ZhuJiahui506
'''

import tensorflow as tf

from main.nn_models import bidirectional_lstm
from main.data_preparation import generate_batch

def sentiment_classify_train(train_texts, train_tags, embedding_filename):
    '''
    进行文本情感分类模型训练 采用堆叠双向LSTM
    :param train_texts: 文本集合 (2d str list)
    :param train_tags: 情感类别标签集合 (numpy 2d array) 每个类别标签采用one-hot表示为向量
    :param embedding_filename: 词向量文件
    '''
    
    batch_size = 128  # mini-batch大小
    padding_size = 250  # 固定每个文本的长度(词汇数) 多截少补
    embedding_size = 200  # 词向量维数
    
    hidden_num = 256  # 隐藏层数目
    layer_num = 3  # 双向LSTM(空间)上层数
    class_num = train_tags.shape[1]  # 类别数目
    
    learning_rate = 0.001  # SGD学习率
    max_iter_num = 20000  # 最大迭代常数
    display_step = 10  # 信息显示间隔
    
    # 定义模型输入数据
    x = tf.placeholder(tf.float32, [None, padding_size, embedding_size], name='x_data')
    y = tf.placeholder(tf.float32, [None, class_num], name='y_data')
    
    # 初始化权重和偏置
    out_weight = tf.Variable(tf.random_normal([2 * hidden_num, class_num]), name='out_weight')
    out_bias = tf.Variable(tf.random_normal([class_num]), name='out_bias')
    
    # 双向LSTM层输出
    prediction = bidirectional_lstm(x, hidden_num, padding_size, embedding_size, out_weight, out_bias, layer_num)
    # print(prediction.name)  # add:0
    
    # 损失函数
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=y))
    # print(loss.name)  # Mean:0
    
    # 优化
    optimizer = tf.train.AdamOptimizer(learning_rate).minimize(loss)
    # print(optimizer.name)  # Adam
    
    # 计算准确率
    correct_pred = tf.equal(tf.argmax(prediction, 1), tf.argmax(y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32), name='accuracy')
    
    # 模型训练
    saver = tf.train.Saver()
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        step = 1
        while step * batch_size < max_iter_num:
            # 产生分片数据
            # x_batch: numpy 3d array, [batch_size, padding_size, embedding_size]
            # y_batch: numpy 2d array, [batch_size, class_num]
            x_batch, y_batch = generate_batch(train_texts, train_tags, batch_size, padding_size, embedding_size, embedding_filename)
            
            # 运行优化函数
            sess.run(optimizer, feed_dict={x: x_batch, y: y_batch})
            
            # 显示信息
            if step % display_step == 0:
                # 得到准确率
                this_accuracy = sess.run(accuracy, feed_dict={x: x_batch, y: y_batch})
                # 得到损失函数值
                this_loss = sess.run(loss, feed_dict={x: x_batch, y: y_batch})
                print("Iter=" + str(step) + ", Minibatch Loss=" + "{:.5f}".format(this_loss) + 
                      ", Training Accuracy=" + "{:.5f}".format(this_accuracy))
            
            step += 1
        
        test_x_batch, test_y_batch = generate_batch(train_texts, train_tags, 500, padding_size, embedding_size, embedding_filename)
        this_accuracy2 = sess.run(accuracy, feed_dict={x: test_x_batch, y: test_y_batch})
        print('Accuracy2', this_accuracy2)
        print('Training finished!')
        # 持久化模型
        save_path = saver.save(sess, 'model/train_model.ckpt')
        print('Save model in ', save_path)


def sentiment_classify_test(test_texts, test_tags):
    '''
    模型测试
    :param test_texts: 文本集合 (2d str list)
    :param test_tags: 情感类别标签集合 (numpy 2d array) 每个类别标签采用one-hot表示为向量
    :return: 测试准确率
    '''

    saver = tf.train.import_meta_graph('model/train_model.ckpt.meta')
    graph = tf.get_default_graph()
    x = graph.get_tensor_by_name('x_data:0')
    y = graph.get_tensor_by_name('y_data:0')
    accuracy = graph.get_tensor_by_name('accuracy:0')
    
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        saver.restore(sess, "model/train_model.ckpt")
        this_accuracy = sess.run(accuracy, feed_dict={x: test_texts, y: test_tags})
    
    return this_accuracy


def sentiment_classify_predict(test_texts, test_tags):
    '''
    模型预测
    :param test_texts: 文本集合 (2d str list)
    :param test_tags: 情感类别标签集合 (numpy 2d array) 每个类别标签采用one-hot表示为向量
    :return: 预测类别标签索引 (numpy 1d int array)
    '''
    
    saver = tf.train.import_meta_graph('model/train_model.ckpt.meta')
    graph = tf.get_default_graph()
    x = graph.get_tensor_by_name('x_data:0')
    prediction = graph.get_tensor_by_name('add:0')
    
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        saver.restore(sess, "model/train_model.ckpt")
        this_y = sess.run(prediction, feed_dict={x: test_texts})
        class_index = tf.argmax(this_y, 1).eval()
    
    return class_index

if __name__ == '__main__':
    pass
